from scipy.spatial.transform import Rotation as R
from matchteris.env import Assemble_Suction
from dm_control.mjcf.element import RootElement
import numpy as np

class GetEnvData:
    '''
        用于获取仿真环境中的数据
        block_ids: 方块颜色列表
        block_height_offset: 方块的高度位置与实际的偏差
    '''

    block_ids = ['blue', 'brown', 'green', 'orange', 'purple', 'red', 'yellow']
    block_height_offset = 0.05

    # 获取吸盘模型
    @classmethod
    def GetSuckerModel(cls, env:Assemble_Suction) -> RootElement:
        return env._vaccum.mjcf_root

    # 获取吸盘位置
    @classmethod
    def GetSuckerPosition(cls, env:Assemble_Suction) -> np.ndarray[float]:
        '''
        获取吸盘的位置的引用，会随着仿真环境的变化而变化
        return: (x, y, z)
        '''
        return env.data.site(env.vaccum_site_id).xpos
    
    # 获取方块id
    @classmethod
    def GetBlockId(cls, env:Assemble_Suction, block_name:str) -> int:
        '''
        获取方块的id
        block_name: color + index + "/", 如"orange0/"
        return: id
        '''
        return env.model.body(block_name).id

    # 获取方块位置
    @classmethod
    def GetBlockPosition(cls, env:Assemble_Suction, block_name:str) -> np.ndarray[float]:
        '''
        获取方块的位置的引用，会随着仿真环境的变化而变化
        block_name: color + index + "/", 如"orange0/"
        return: (x, y, z)
        '''
        block_id = env.model.body(block_name).id
        block_xpos = env.data.xpos[block_id]
        return block_xpos
    
    # 获取吸盘转过的角度
    @classmethod
    def GetSuckerAngle(cls, env:Assemble_Suction) -> float:
        sucker_angle = env.data.ctrl[-2]
        sucker_angle = np.rad2deg(sucker_angle)
        return sucker_angle